#!/usr/bin/env python
# -*- coding: utf-8 -*-
#只是识别绿色的无人机测试

from traceback import print_tb
import numpy as np
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from geometry_msgs.msg import Twist
import argparse
import cv2
import sys
from djitellopy import Tello
import math
# import cv2.aruco

tello = Tello()
class RAS_Tello(Node):
    def __init__(self):
        super().__init__('tello_project')

        self.subscription = self.create_subscription(
            Image,
            '/image_raw',
            self.image_callback,
            10
        )

        self.publisher_ = self.create_publisher(Twist, '/control', 10)
        self.tello = Tello()

    def image_callback(self, msg):
        
        self.frame = self.imgmsg_to_cv2(msg)   ###-------------------------------------------------
        # RGB to HSV color space
        hsv = cv2.cvtColor(self.frame, cv2.COLOR_BGR2HSV)
        # structural element
        line = cv2.getStructuringElement(cv2.MORPH_RECT, (15, 15), (-1, -1))
        # HSV
        mask_white = cv2.inRange(hsv, (0,0,221), (180,30,255))
        # mask_blue = cv2.inRange(hsv, (100, 43, 46), (124,255,255))
        mask_green = cv2.inRange(hsv, (35,43,46), (77,255,255))
        # mask_yellow = cv2.inRange(hsv,(26,43,46),(34,255,255))
        mask = mask_green
        #Contour extraction, find the largest contour 
        i = 1
        # for mask in masks:
    
        contours, hierarchy = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        index = -1
        max = 0
        for c in range(len(contours)):
            area = cv2.contourArea(contours[c])
            if area > max:
                max = area
                index = c
        # pic
        if index >= 0:
            rect = cv2.minAreaRect(contours[index])
            #Ellipse Fitting
            cv2.ellipse(self.frame, rect, (0, 255, 0), 2, 8)
            #center point positioning
            cv2.circle(self.frame, (np.int32(rect[0][0]), np.int32(rect[0][1])), 2, (0, 255, 0), 2, 8, 0)
        if i == 1:
            text = 'green'
        # elif i == 2:
        #     text = 'blue'
        # elif i == 3:
        #     text = 'green'
        # elif i == 4:
        #     text = 'yellow'
        # cv2.putText(self.frame,text,(np.int32(rect[0][0]), np.int32(rect[0][1]) ), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 2)
        # i += 1
        cv2.imshow("Frame", self.frame)
        cv2.waitKey(1)
        

    def tello_movement(self):
        msg = Twist()
        print("[INFO] moving Tello...")
        self.publisher_.publish(msg)
        print('[INFO] Tello moved')

    def imgmsg_to_cv2(self, img_msg):
        n_channels = len(img_msg.data) // (img_msg.height * img_msg.width)
        dtype = np.uint8

        img_buf = np.asarray(img_msg.data, dtype=dtype) if isinstance(img_msg.data, list) else img_msg.data

        if n_channels == 1:
            cv2_img = np.ndarray(shape=(img_msg.height, img_msg.width), dtype=dtype, buffer=img_buf)
        else:
            cv2_img = np.ndarray(shape=(img_msg.height, img_msg.width, n_channels), dtype=dtype, buffer=img_buf)

        # If the byte order is different between the message and the system
        if img_msg.is_bigendian == (sys.byteorder == 'little'):
            cv2_img = cv2_img.byteswap().newbyteorder()

        return cv2_img

def main(args=None):
    rclpy.init(args=args)
    tello_project = RAS_Tello()
    rclpy.spin(tello_project)
    tello_project.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()